import pygame
import random

# 初始化 Pygame
pygame.init()

# 游戏区域大小
WIDTH, HEIGHT = 400, 400

# 网格大小和数量
GRID_SIZE = 20
GRID_WIDTH, GRID_HEIGHT = WIDTH // GRID_SIZE, HEIGHT // GRID_SIZE

# 设置屏幕尺寸和标题
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('Snake Game')

# 定义颜色
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)

# 定义字体
FONT = pygame.font.SysFont(None, 30)

# 定义蛇的初始位置和长度
snake_init_pos = [(GRID_WIDTH // 2, GRID_HEIGHT // 2)]
snake_init_length = 3

# 定义贪吃蛇的移动方向
UP = 0
RIGHT = 1
DOWN = 2
LEFT = 3

# 初始化贪吃蛇和食物
snake = snake_init_pos.copy()
direction = RIGHT
food = []

# 随机生成食物位置
def generate_food():
    x = random.randint(0, GRID_WIDTH - 1)
    y = random.randint(0, GRID_HEIGHT - 1)
    food.append((x, y))

generate_food()

# 移动贪吃蛇
def move_snake():
    global snake, direction
    
    head_x, head_y = snake[-1]
    
    if direction == UP:
        head_y -= 1
    elif direction == DOWN:
        head_y += 1
    elif direction == LEFT:
        head_x -= 1
    elif direction == RIGHT:
        head_x += 1
    
    snake.append((head_x, head_y))
    
    # 移除尾部
    if len(snake) > snake_init_length:
        del snake[0]

# 处理事件
def handle_events():
    global direction
    
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            quit()
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and direction != DOWN:
                direction = UP
            elif event.key == pygame.K_RIGHT and direction != LEFT:
                direction = RIGHT
            elif event.key == pygame.K_DOWN and direction != UP:
                direction = DOWN
            elif event.key == pygame.K_LEFT and direction != RIGHT:
                direction = LEFT

# 碰撞检测
def collision_detection():
    global snake_init_length, food
    
    head = snake[-1]
    
    # 检测是否吃到食物
    if head in food:
        generate_food()
        snake_init_length += 1
        food.remove(head)
    
    # 检测是否撞墙
    if head[0] < 0 or head[0] >= GRID_WIDTH or \
       head[1] < 0 or head[1] >= GRID_HEIGHT:
        return True
    
    # 检测是否撞到自己的身体
    if head in snake[:-1]:
        return True
    
    return False

# 绘制网格线
def draw_grid():
    for x in range(0, WIDTH, GRID_SIZE):
        pygame.draw.line(screen, BLACK, (x, 0), (x, HEIGHT))
    for y in range(0, HEIGHT, GRID_SIZE):
        pygame.draw.line(screen, BLACK, (0, y), (WIDTH, y))

# 绘制贪吃蛇和食物
def draw_snake_and_food():
    # 绘制贪吃蛇的身体
    for segment in snake:
        pygame.draw.rect(screen, GREEN, (segment[0] * GRID_SIZE, segment[1] * GRID_SIZE, GRID_SIZE, GRID_SIZE))
    
    # 绘制食物
    for f in food:
        pygame.draw.circle(screen, RED, (f[0] * GRID_SIZE + GRID_SIZE // 2, f[1] * GRID_SIZE + GRID_SIZE // 2), GRID_SIZE // 2)

# 绘制得分
def draw_score(score):
    score_text = FONT.render('Score: {}'.format(score), True, BLACK)
    screen.blit(score_text, (10, 10))

# 初始化得分
score = 0

# 游戏循环
while True:
    # 处理事件
    handle_events()
    
    # 移动贪吃蛇
    move_snake()
    
    # 碰撞检测
    if collision_detection():
        break
    
    # 绘制背景和网格线
    screen.fill(WHITE)
    draw_grid()
    
    # 绘制贪吃蛇和食物
    draw_snake_and_food()
    
    # 绘制得分
    draw_score(score)
    
    # 刷新屏幕
    pygame.display.update()
    
    # 控制帧率
    pygame.time.Clock().tick(10)
    
# 游戏结束，显示得分并等待关闭窗口
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            quit()
    
    # 绘制背景和网格线
    screen.fill(WHITE)
    draw_grid()
    
    # 显示得分
    score_text = FONT.render('Game over! Your score: {}'.format(score), True, BLACK)
    screen.blit(score_text, (WIDTH // 2 - score_text.get_width() // 2, HEIGHT // 2 - score_text.get_height() // 2))
    
    # 刷新屏幕
    pygame.display.update()
